# -*- coding:utf-8 -*-

from airflow.hooks.base_hook import BaseHook
from sqlalchemy import create_engine
from sqlalchemy.orm import scoped_session, sessionmaker
import setting


class DBHook(BaseHook):

    def Session(self, conn_id):
        SQL_ALCHEMY_CONN = self.get_uri(conn_id)
        print SQL_ALCHEMY_CONN
        # engine = create_engine(SQL_ALCHEMY_CONN,echo=True)
        engine = create_engine(SQL_ALCHEMY_CONN,echo = setting.echo)
        Session = scoped_session(
            sessionmaker(autocommit=False, autoflush=False, bind=engine))
        return Session

    def get_uri(self, conn_id):
        conn = self.get_connection(conn_id)

        login = ''
        if conn.login:
            login = '{conn.login}:{conn.password}@'.format(conn=conn)
        host = conn.host
        if conn.port is not None:
            host += ':{port}'.format(port=conn.port)
        ext = ''
        if conn.extra is not None:
            s = conn.extra_dejson

            ext = '?charset={}'.format(s['charset'])

        return '{conn.conn_type}://{login}{host}/{conn.schema}{ext}'.format(
            conn=conn, login=login, host=host, ext=ext)

    def Connection(self,conn_id):
        SQL_ALCHEMY_CONN = self.get_uri(conn_id)
        engine = create_engine(SQL_ALCHEMY_CONN,echo = setting.echo)
        conn = engine.connect()
        return conn

db_hook = DBHook('db_hook')